#!/usr/bin/python

# 标准库导入
import os
import logging
import threading
from typing import Optional, Tuple, Dict, Any, List
from concurrent.futures import ThreadPoolExecutor, as_completed

# 第三方库导入
import yaml
import paramiko
import zipfile

# 日志配置
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('logs/deploy.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)


class SSHConnection:
    """SSH连接管理类，处理远程服务器连接和文件传输
    
    支持密码和证书两种认证方式，提供文件传输和远程命令执行功能。
    实现了上下文管理器接口，可以使用with语句自动管理连接的建立和关闭。
    
    Attributes:
        hostname: 服务器主机名
        port: SSH端口号
        username: 登录用户名
        auth_type: 认证类型，支持'密码'和'证书'
        password: SSH密码（密码认证时使用）
        key_path: PEM证书路径（证书认证时使用）
        client: SSH客户端连接实例
        sftp: SFTP客户端连接实例
    """
    
    def __init__(self, hostname: str, port: int, username: str, auth_type: str = '密码', password: str = None, key_path: str = None) -> None:
        """初始化SSH连接实例
        
        Args:
            hostname: 服务器主机名
            port: SSH端口号
            username: 登录用户名
            auth_type: 认证类型，支持'密码'和'证书'，默认为'密码'
            password: SSH密码（密码认证时使用）
            key_path: PEM证书路径（证书认证时使用）
        """
        self.hostname = hostname
        self.port = port
        self.username = username
        self.auth_type = auth_type
        self.password = password
        self.key_path = key_path
        self.client: Optional[paramiko.SSHClient] = None
        self.sftp: Optional[paramiko.SFTPClient] = None

    def __enter__(self) -> 'SSHConnection':
        self.connect()
        return self

    def __exit__(self, exc_type, exc_value, traceback) -> None:
        self.close()

    def connect(self) -> None:
        """建立SSH连接，支持密码和证书两种认证方式"""
        self.client = paramiko.SSHClient()
        self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        
        try:
            if self.auth_type == '密码':
                self.client.connect(
                    self.hostname,
                    port=self.port,
                    username=self.username,
                    password=self.password
                )
            else:
                if not os.path.exists(self.key_path):
                    raise FileNotFoundError(f'PEM文件不存在: {self.key_path}')
                
                # 检查文件权限（仅在类Unix系统上）
                if os.name != 'nt':
                    import stat
                    current_mode = os.stat(self.key_path).st_mode
                    if current_mode & (stat.S_IRWXG | stat.S_IRWXO):
                        logger.warning('PEM文件权限过于开放，AWS要求权限为400')
                        # 自动修复权限
                        os.chmod(self.key_path, stat.S_IRUSR | stat.S_IWUSR)
                
                try:
                    pkey = paramiko.RSAKey.from_private_key_file(self.key_path)
                except paramiko.ssh_exception.PasswordRequiredException:
                    raise ValueError('PEM文件受密码保护，请提供解密密码')
                except paramiko.ssh_exception.SSHException as e:
                    raise ValueError(f'PEM文件格式错误: {e}')
                
                try:
                    self.client.connect(
                        self.hostname,
                        port=self.port,
                        username=self.username,
                        pkey=pkey
                    )
                except paramiko.ssh_exception.AuthenticationException:
                    raise ValueError('认证失败，请检查用户名和PEM文件是否匹配')
                except paramiko.ssh_exception.NoValidConnectionsError as e:
                    raise ConnectionError(f'无法连接到服务器: {e}')
                
            self.sftp = self.client.open_sftp()
        except Exception as e:
            logger.error(f'SSH连接失败: {str(e)}')
            raise

    def close(self) -> None:
        """关闭SSH连接"""
        if self.sftp:
            self.sftp.close()
        if self.client:
            self.client.close()

    def execute_command(self, command: str) -> Tuple[str, str]:
        """执行远程命令"""
        if not self.client:
            raise RuntimeError("SSH连接未建立")
        stdin, stdout, stderr = self.client.exec_command(command)
        return stdout.read().decode('utf-8'), stderr.read().decode('utf-8')

    def upload_file(self, local_path: str, remote_path: str) -> None:
        """上传文件到远程服务器"""
        if not self.sftp:
            raise RuntimeError("SFTP连接未建立")
        self.sftp.put(local_path, remote_path)

    def change_directory(self, path: str) -> None:
        """切换远程目录"""
        if not self.sftp:
            raise RuntimeError("SFTP连接未建立")
        self.sftp.chdir(path)


class FileManager:
    """文件操作管理类，处理本地文件压缩和命令执行
    
    提供静态方法用于处理文件压缩和本地命令执行操作。
    所有方法都包含完整的错误处理和日志记录。
    """

    @staticmethod
    def compress_directory(directory_path: str, output_path: str) -> None:
        """压缩目录到zip文件"""
        try:
            with zipfile.ZipFile(output_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
                for root, _, files in os.walk(directory_path):
                    for file in files:
                        file_path = os.path.join(root, file)
                        zip_path = os.path.relpath(file_path, directory_path)
                        zipf.write(file_path, arcname=os.path.join(os.path.basename(directory_path), zip_path))
            logger.info(f"成功压缩目录 {directory_path} 到 {output_path}")
        except Exception as e:
            logger.error(f"压缩目录失败: {e}")
            raise

    @staticmethod
    def execute_local_command(command: str) -> str:
        """执行本地命令"""
        import subprocess
        try:
            result = subprocess.run(command, shell=True, text=True, encoding='utf-8', capture_output=True)
            if result.returncode != 0:
                logger.error(f"执行命令失败: {result}")
                raise RuntimeError(f"命令执行失败: {command} {result.stderr}")
            logger.info(f"成功执行命令: {command}")
            return result.stdout
        except Exception as e:
            logger.error(f"执行命令失败: {e}")
            raise


class Project:
    """项目类，统一处理项目的构建和部署流程
    
    负责单个项目的构建和部署流程管理，包括：
    1. 执行本地构建命令
    2. 处理构建产物（可选压缩）
    3. 上传构建产物到远程服务器
    4. 执行远程部署脚本
    
    Attributes:
        config: 项目配置信息
        file_manager: 文件管理器实例
        ssh_connection: SSH连接实例
        project_name: 项目名称
        enabled: 项目是否启用
        local_dir: 本地项目目录
        remote_dir: 远程部署目录
        commands: 构建命令列表
        deploy_script: 部署脚本路径
        artifacts: 构建产物配置
    """
    
    # 类级别的目录锁，用于确保多线程执行时的目录切换安全
    _dir_lock = threading.Lock()

    def __init__(self, project_config: Dict[str, Any], file_manager: FileManager, ssh_connection: SSHConnection) -> None:
        self.config = project_config
        self.file_manager = file_manager
        self.ssh_connection = ssh_connection
        self.project_name = project_config['name']
        self.enabled = project_config.get('enabled', True)  # 默认启用
        self.local_dir = project_config['local_dir']
        self.remote_dir = project_config['remote_dir']
        self.commands = project_config.get('commands', [])
        self.deploy_script = project_config.get('deploy_script')
        self.artifacts = project_config.get('artifacts', {})

    def build(self) -> None:
        """构建项目"""
        try:
            logger.info(f'开始构建{self.project_name}项目...')
            # 使用目录锁确保线程安全的目录切换
            with self._dir_lock:
                # 保存当前工作目录
                original_dir = os.getcwd()
                try:
                    os.chdir(self.local_dir)

                    # 执行构建命令
                    for command in self.commands:
                        self.file_manager.execute_local_command(command)

                    # 处理构建产物
                    if self.artifacts.get('compress'):
                        # 确保路径使用Unix格式的正斜杠
                        source_dir_path = self.artifacts['source_dir'].replace('\\', '/')
                        source_dir = os.path.join(self.local_dir, source_dir_path)
                        output_file = self.artifacts['output_file'].replace('\\', '/')
                        self.file_manager.compress_directory(source_dir, output_file)

                    logger.info(f'{self.project_name}项目构建完成')
                finally:
                    # 恢复原始工作目录
                    os.chdir(original_dir)
        except Exception as e:
            logger.error(f'{self.project_name}项目构建失败: {e}')
            raise

    def deploy(self) -> None:
        """部署项目"""
        try:
            logger.info(f'开始部署{self.project_name}项目...')
            # 使用目录锁确保线程安全的目录切换
            with self._dir_lock:
                self.ssh_connection.change_directory(self.remote_dir)

                # 上传构建产物
                if self.artifacts:
                    # 确保路径使用Unix格式的正斜杠
                    output_file = self.artifacts['output_file'].replace('\\', '/')
                    local_artifact = os.path.join(self.local_dir, output_file)
                    remote_artifact = os.path.basename(output_file)
                    self.ssh_connection.upload_file(local_artifact, remote_artifact)

                # 执行部署脚本
                if self.deploy_script:
                    command = f'/bin/bash {self.deploy_script}'
                    stdout, stderr = self.ssh_connection.execute_command(command)
                    if stderr:
                        raise RuntimeError(f'{self.project_name}部署命令执行失败: {stderr}')
                    logger.info(f'执行{self.project_name}部署命令: {command}\n{stdout}')

                logger.info(f'{self.project_name}项目部署完成')
        except Exception as e:
            logger.error(f'{self.project_name}项目部署失败: {e}')
            raise


class DeploymentManager:
    """部署管理器，负责协调项目的构建和部署
    
    作为整个部署流程的核心协调器，主要职责包括：
    1. 加载和解析配置文件
    2. 管理SSH连接的生命周期
    3. 初始化和管理所有项目实例
    4. 协调执行所有项目的构建和部署流程
    
    实现了上下文管理器接口，可以使用with语句自动管理资源。
    """

    def __init__(self, config_file: str = 'config.yaml') -> None:
        self.config = self._load_config(config_file)
        server_config = self.config['server']
        self.ssh_connection = SSHConnection(
            hostname=server_config['hostname'],
            port=server_config['port'],
            username=server_config['username'],
            auth_type=server_config.get('auth_type', '密码'),
            password=server_config.get('password'),
            key_path=server_config.get('key_path')
        )
        self.file_manager = FileManager()
        self.projects = self._init_projects()

    def __enter__(self) -> 'DeploymentManager':
        self.ssh_connection.__enter__()
        return self

    def __exit__(self, exc_type, exc_value, traceback) -> None:
        self.ssh_connection.__exit__(exc_type, exc_value, traceback)

    def _load_config(self, config_file: str) -> Dict[str, Any]:
        """加载配置文件"""
        try:
            with open(config_file, 'r', encoding='utf-8') as f:
                return yaml.safe_load(f)
        except Exception as e:
            logger.error(f"加载配置文件失败: {e}")
            raise

    def _init_projects(self) -> List[Project]:
        """初始化项目列表"""
        projects = []
        for project_config in self.config['projects']:
            project = Project(project_config, self.file_manager, self.ssh_connection)
            projects.append(project)
        return projects

    def _deploy_project(self, project: Project) -> None:
        """执行单个项目的构建和部署"""
        try:
            if not project.enabled:
                logger.info(f'跳过未启用的项目: {project.project_name}')
                return
            project.build()
            project.deploy()
            logger.info(f'{project.project_name}项目部署完成')
        except Exception as e:
            logger.error(f'{project.project_name}项目部署失败: {e}')
            raise

    def deploy_all(self) -> None:
        """使用多线程并行执行完整的构建部署流程"""
        try:
            with ThreadPoolExecutor(max_workers=min(len(self.projects), 4)) as executor:
                # 提交所有项目的部署任务到线程池
                future_to_project = {executor.submit(self._deploy_project, project): project 
                                   for project in self.projects}
                
                # 等待所有任务完成并处理结果
                for future in as_completed(future_to_project):
                    project = future_to_project[future]
                    try:
                        future.result()
                    except Exception as e:
                        logger.error(f'{project.project_name}项目部署失败: {e}')
                        raise
                        
            logger.info('所有项目部署流程全部完成')
        except Exception as e:
            logger.error(f'部署过程发生错误: {e}')
            raise


if __name__ == '__main__':
    try:
        with DeploymentManager('conf/config.yaml') as manager:
            manager.deploy_all()
        logger.info('部署完成')
    except Exception as e:
        logger.error(f'部署失败: {e}')
        exit(1)
